#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <spu_intrinsics.h>
#include <spu_mfcio.h>

#include "fish.h"
#include "flocking.h"

#define VERSION 7

/* Flocking Parameters */
const float C0 = 0.00000001; // weight of vector away from neighbors within crowding radius
//const float C1 = 2000.0; // weight of vector towards center of world
const float C1 = 0.0;
const float C2 = 1.0; // weight of velocity correcting vectors towards avg. flocking velocity
const float C3 = 0.0; // weight of vector towards center of gravity
const float C4 = 0.05; // weight of random vector
const float C5 = 0.1; // weight of vector pointing to cog of local flockmates
const float R_CROWD2 = 0.00390625; // square of crowding radius
const float R_SIGHT2 = 0.015625; // square of sighting radius
const float R_SIGHT = 0.125; // sighting radius
const float BIRD_MASS = 1.0; // bird mass
const float R = 1.0; // repulsion constant
const float M = 1.0; // constant that converts vectors to acceleration
const float START_CONST = 0.5;
const float SIGHT_ANGLE = 200.0 * M_PI / 180.0;
const float SCALE_FACTOR = 1.0;
float v0x, v0y, v1x, v1y, v2x, v2y, v3x, v3y, v4x, v4y, v5x, v5y, vx, vy, ax, ay; // vectors
float cogx, cogy; // center of gravity coordinates
float local_cogx, local_cogy;
float tx, ty; // translation vectors for the moving grid

fish_t* fish[BUCKETS_Y][BUCKETS_X];
int fish_bucket_n[2][BUCKETS_Y][BUCKETS_X] __attribute__((aligned(16)));
int n_fish;

struct Timer dma_timer;
struct Timer mailbox_timer;
struct Timer computation_timer;

void interact_fish_bucket_with_bucket(int, int, int, int, int);
void move_fish(int, float, int, int);
void bounce_fish (fish_t*);

int main(unsigned long long spe, unsigned long long argp, unsigned long long envp)
{
  // DMA in the spe_params_t struct
  spe_params_t spe_params __attribute__((aligned(16)));
  spe_auxinfo_t spe_auxinfo __attribute__((aligned(16)));
  float dt;
  int tag, t, bucket_x, bucket_y, offset_x, offset_y, bucket_number, n, m, n_prev, m_prev, n0, m0;
  unsigned int seed;
  char virgin_flag;
  spe_mbox_msg_t spe_status_msg, spe_action_msg, spe_dt_msg;

  initialize_timer(&dma_timer);
  initialize_timer(&mailbox_timer);
  initialize_timer(&computation_timer);


  //printf("(Z, Z): Initiating Parameter Transfer\n\tmfc_ea2h(argp): %lu\n\tmfc_ea2l(argp): %lu\n\tmfc_ea2h(envp): %lu\n\tmfc_ea2l(envp): %lu\n", mfc_ea2h(argp), mfc_ea2l(argp), mfc_ea2h(envp), mfc_ea2l(envp));

  // issue the transfer
  tag = 1;
  spu_mfcdma64(&spe_params, mfc_ea2h(argp), mfc_ea2l(argp), sizeof(spe_params_t), tag, MFC_GET_CMD);
  spu_mfcdma64(&spe_auxinfo, mfc_ea2h(envp), mfc_ea2l(envp), sizeof(spe_auxinfo_t), tag, MFC_GET_CMD);

  
  start_timer(&dma_timer);
  // wait for all data with tag == 1
  spu_writech(MFC_WrTagMask, 1 << 1);
  spu_mfcstat(MFC_TAG_UPDATE_ALL);
  stop_timer(&dma_timer);

  bucket_number = spe_auxinfo.bucket_number;
  seed = spe_auxinfo.random_seed*((unsigned int)(-1));
  n = bucket_number / BUCKETS_X;
  m = bucket_number % BUCKETS_X;
  srand48(seed); 
  
  /*  
  printf("(%d, %d): Parameter Transfer Completed. Parameters:\n\tbucket#: %d\n\tfish_ea[0] array:\n\t\t%lu\t%lu\t%lu\n\t\t%lu\t%lu\t%lu\n\tfish_ea[1] array:\n\t\t%lu\t%lu\t%lu\n\t\t%lu\t%lu\t%lu\n\tmfc_ea2h(fish_bucket_n_ea):\t%lu\n\tmfc_ea2l(fish_bucket_n_ea):\t%lu\n", 
	 n, m, bucket_number, 
	 mfc_ea2l(spe_params.fish_ea[0][0][0]),
	 mfc_ea2l(spe_params.fish_ea[0][0][1]),
	 mfc_ea2l(spe_params.fish_ea[0][0][2]),
	 mfc_ea2l(spe_params.fish_ea[0][1][0]),
	 mfc_ea2l(spe_params.fish_ea[0][1][1]),
	 mfc_ea2l(spe_params.fish_ea[0][1][2]),
	 mfc_ea2l(spe_params.fish_ea[1][0][0]),
	 mfc_ea2l(spe_params.fish_ea[1][0][1]),
	 mfc_ea2l(spe_params.fish_ea[1][0][2]),
	 mfc_ea2l(spe_params.fish_ea[1][1][0]),
	 mfc_ea2l(spe_params.fish_ea[1][1][1]),
	 mfc_ea2l(spe_params.fish_ea[1][1][2]),
	 mfc_ea2h(spe_params.fish_bucket_n_ea),
	 mfc_ea2l(spe_params.fish_bucket_n_ea) );
  */

  // here we'll do our main loop
  // on each iteration of the loop we need to
  // 1) get the go or stop message from ppe
  // 2) if go msg then:
  // 3) DMA in the fish_bucket_n array
  // 4) Use this to DMA in the fish
  // 5) do normal thing
  // 6) if stop msg received then exit the code

  t = 0;

  start_timer(&mailbox_timer);
  spe_action_msg.i = spu_read_in_mbox(); // get the very first action msg
  stop_timer(&mailbox_timer);

  while( spe_action_msg.i == GO_MSG ) {    
    // DMA in the fish_bucket_n
    tag = 1;
    spu_mfcdma64(fish_bucket_n, mfc_ea2h((unsigned long long)spe_params.fish_bucket_n_ea), mfc_ea2l((unsigned long long)spe_params.fish_bucket_n_ea), 
		 2*BUCKETS_X*BUCKETS_Y*sizeof(int), tag, MFC_GET_CMD);

    start_timer(&mailbox_timer);
    // overlap getting dt from mailbox w/ DMA transfer
    spe_dt_msg.i = spu_read_in_mbox();
    stop_timer(&mailbox_timer);
    dt = spe_dt_msg.f; // wahahaha

    start_timer(&dma_timer);
    // wait for all data with tag == 1
    spu_writech(MFC_WrTagMask, 1 << 1);
    spu_mfcstat(MFC_TAG_UPDATE_ALL);
    stop_timer(&dma_timer);

        /*
    printf("(%d, %d): fish_bucket_n Transfer Completed.\n\tfish_bucket_n[0] array:\n\t\t%d\t%d\t%d\n\t\t%d\t%d\t%d\n\tfish_bucket_n[1] array:\n\t\t%d\t%d\t%d\n\t\t%d\t%d\t%d\n", 
	   n, m, 
	   fish_bucket_n[0][0][0],
	   fish_bucket_n[0][0][1],
	   fish_bucket_n[0][0][2],
	   fish_bucket_n[0][1][0],
	   fish_bucket_n[0][1][1],
	   fish_bucket_n[0][1][2],
	   fish_bucket_n[1][0][0],
	   fish_bucket_n[1][0][1],
	   fish_bucket_n[1][0][2],
	   fish_bucket_n[1][1][0],
	   fish_bucket_n[1][1][1],
	   fish_bucket_n[1][1][2] );
    */

    fish[n][m] = (fish_t*)memalign(16, (fish_bucket_n[t][n][m]+1)*sizeof(fish_t));

    // firstly, DMA in our own fish and our left neighbor, to get us started
    tag = 0;
    spu_mfcdma64(fish[n][m], mfc_ea2h((unsigned long long)spe_params.fish_ea[t][n][m]), mfc_ea2l((unsigned long long)spe_params.fish_ea[t][n][m]), 
		 fish_bucket_n[t][n][m]*sizeof(fish_t), tag, MFC_GET_CMD);

    // for diabolical multibuffering plan of attack
    n_prev = n;
    m_prev = m;

    // DMA in the rest of the fish and overlap with computation
    for( offset_y = 0; offset_y < 2; offset_y++ ) {
      for( offset_x = -1; offset_x < 2; offset_x++ ) {
	if( !(offset_y == 0 && offset_x == 0) ) { // else, skip
	 
	  tag ^= 1; // oscillate the tag
	  n0 = (n+offset_y+BUCKETS_Y) % BUCKETS_Y;
	  m0 = (m+offset_x+BUCKETS_X) % BUCKETS_X;    
	  
	  fish[n0][m0] = (fish_t*)memalign(16, (fish_bucket_n[t][n0][m0]+1)*sizeof(fish_t));
	  
	  // issue DMA for the current bucket (COME BACK AND TACKLE ZERO FISH IN BUCKET CASE!!!)
	  //printf("(%d, %d): Transfering %d fish from bucket (%d, %d)\n", n, m, fish_bucket_n[t][n0][m0], n0, m0);
	  spu_mfcdma64(fish[n0][m0], mfc_ea2h((unsigned long long)spe_params.fish_ea[t][n0][m0]), mfc_ea2l((unsigned long long)spe_params.fish_ea[t][n0][m0]), 
		       fish_bucket_n[t][n0][m0]*sizeof(fish_t), tag, MFC_GET_CMD);
	  
	  start_timer(&dma_timer);
	  // wait for previous iteration's DMA to finish
	  spu_writech(MFC_WrTagMask, 1 << (tag^1));
	  spu_mfcstat(MFC_TAG_UPDATE_ALL);
	  stop_timer(&dma_timer);

	  start_timer(&computation_timer);
	  interact_fish_bucket_with_bucket(t, n, m, n_prev, m_prev);
	  stop_timer(&computation_timer);

	  n_prev = n0;
	  m_prev = m0;
	}
      }
    }

    start_timer(&dma_timer);
    // don't forget the very last guy we DMA'd but haven't interacted with
    spu_writech(MFC_WrTagMask, 1 << tag);
    spu_mfcstat(MFC_TAG_UPDATE_ALL);   

    stop_timer(&dma_timer);

    start_timer(&computation_timer);
    interact_fish_bucket_with_bucket(t, n, m, n_prev, m_prev);

    

    // do a simple test to see if we grabbed from an old place in memory
    //if( (fish[n][m][0].x == 0.0) && (fish[n][m][0].y == 0.0) && (fish[n][m][0].vx == 0.0) && (fish[n][m][0].vy == 0.0) )
    //printf("(%d, %d): Grabbing data from shadow memory\n", n, m);
    //printf("(%d, %d): t = %d, dt = %f, data src/dst: %lu\n", n, m, t, dt, mfc_ea2l((unsigned long long)spe_params.fish_ea[t][n][m]) );
    
    move_fish(t, dt, n, m);
    stop_timer(&computation_timer);

    // now we need to write the fish back to main memory
    spu_mfcdma64(fish[n][m], mfc_ea2h((unsigned long long)spe_params.fish_ea[t][n][m]), mfc_ea2l((unsigned long long)spe_params.fish_ea[t][n][m]), 
		 fish_bucket_n[t][n][m]*sizeof(fish_t), tag, MFC_PUT_CMD);
    start_timer(&dma_timer);
    spu_writech(MFC_WrTagMask, 1 << 1);
    spu_mfcstat(MFC_TAG_UPDATE_ALL);
    stop_timer(&dma_timer);

    // tell PPE, "I'm done!"
    spe_status_msg.i = DONE_MSG;
    spu_write_out_mbox(spe_status_msg.i); // send our status to PPE

    // free the fish memory, it's useless now
    for( offset_x = -1; offset_x<2; offset_x++ ) {
      for( offset_y = 0; offset_y<2; offset_y++ ) {
	n0 = (n+offset_y+BUCKETS_Y) % BUCKETS_Y;
	m0 = (m+offset_x+BUCKETS_X) % BUCKETS_X;
	free( fish[n0][m0] );
      }
    }
    
    start_timer(&mailbox_timer);
    spe_action_msg.i = spu_read_in_mbox(); // a blocking call
    stop_timer(&mailbox_timer);

    t ^= 1;
  }

  //printf("(%d, %d): Fish Transfer Completed\n", n, m);

  // by this time we have the fish array filled up, and the fish_bucket_n 
  // array filled up... now we can run the interact_fish_bucket function from
  // our old code just fine.

  // calculate all the fish's ax and ay's
  printf("version\t%d\tfish\t%d\tmax_steps\t%d\tbucket_number\t%d\tdma_timer\t%g\tmailbox_timer\t%g\tcomputation_timer\t%g\n", 
	 VERSION, 256, 100, bucket_number, timer_duration(dma_timer), timer_duration(mailbox_timer), timer_duration(computation_timer));
  

  return 0;
}


void
move_fish(int t, float dt, int n, int m)
{
  int i;

  for(i = 0; i < fish_bucket_n[t][n][m]; i++) {
    fish_t *__restrict__ f = &fish[n][m][i];
    
    f->x  += dt * f->vx * SCALE_FACTOR;
    f->y  += dt * f->vy * SCALE_FACTOR;
    f->vx += dt * f->ax;
    f->vy += dt * f->ay;
    
    bounce_fish(f);
  }
}


void
interact_fish_bucket_with_bucket(int t, int n, int m, int target_n, int target_m)
{
  int i, j, ni, mi, n0, m0;
  fish_t *__restrict__ wfish;
  fish_t const *__restrict__ rfish;
  float r2, r16, r32, r4, r, a, angle, vtempx, vtempy, dx, dy;
  int in_sight, in_crowd;

  wfish = (fish_t *__restrict__)fish[n][m]; // our fish
  rfish = (fish_t const *__restrict__)fish[target_n][target_m]; // target fish

  // for every fish in our bucket
  for(i = 0; i < fish_bucket_n[t][n][m]; i++){

    v0x = v0y = v1x = v1y = v2x = v2y = v3x = v3y = v4x = v4y = v5x = v5y = vx = vy = ax = ay = 0.0;
    in_sight = in_crowd = 0;
    local_cogx = 0.0;
    local_cogy = 0.0;

    // calculate random unit vector
    v4x = drand48() - 0.5;
    v4y = drand48() - 0.5;
    r2 = f_max(v4x * v4x + v4y * v4y, min_r2);
    r = sqrt(r2);
    v4x /= r;
    v4y /= r;

    // interact with fish in the target bucket!!!

    for(j = 0; j < fish_bucket_n[t][target_n][target_m]; j++){
      dx = rfish[j].x - wfish[i].x;
      if ( fabs( dx ) > WALL_SEP - R_SIGHT ) {
	if( dx < 0 )
	  dx = fabs( dx ) - WALL_SEP;
	else
	  dx = WALL_SEP - fabs( dx );
      }
      dy = rfish[j].y - wfish[i].y;
      if ( fabs( dy ) > WALL_SEP - R_SIGHT ) {
	if( dy < 0 )
	  dy = fabs( dy ) - WALL_SEP;
	else
	  dy = WALL_SEP - fabs( dy );
      }

      if (dx != 0 || dy != 0) { 
	r2 = f_max(dx * dx + dy * dy, min_r2);

	// figure out if the boid is in this boid's range of vision
	angle = fabs( atan2(dy, dx) - atan2(wfish[i].vy, wfish[i].vx) );

	if( (angle <= (SIGHT_ANGLE/2.0)) || (angle >= (2*M_PI - SIGHT_ANGLE/2.0)) ) {
	  
	  // calculate crowding vector
	  if ( r2 < R_CROWD2 ) {
	    v0x -= dx;
	    v0y -= dy;
	    in_crowd++;
	  }
	  
	  // calculate flocking vector
	  if( r2 < R_SIGHT2 ) {
	    v2x += rfish[j].vx;
	    v2y += rfish[j].vy;
	    
	    local_cogx += rfish[j].x;
	    local_cogy += rfish[j].y;
	    
	    in_sight++;
	  }
	}
      } 
    }

    if( in_sight != 0 ) {
      // calculate vector pointing in the direction of the average velocity of nearby flockmates
      v2x /= in_sight;
      v2y /= in_sight;

      v2x = v2x - wfish[i].vx;
      v2y = v2y - wfish[i].vy;

      local_cogx /= in_sight;
      local_cogy /= in_sight;
      
      // calculate vector towards cog of local flockmates
      v5x = local_cogx - wfish[i].x;
      v5y = local_cogy - wfish[i].y;
    }

    if( in_crowd != 0 ) {
      v0x /= in_crowd;
      v0y /= in_crowd;
      
      r2 = v0x * v0x + v0y * v0y;

      // let the repulsion decrease with r4
      v0x /= r2*r2;
      v0y /= r2*r2;
    }

    // now we have v0 to v5...
    vx = C0*v0x + C1*v1x + C2*v2x + C3*v3x + C4*v4x + C5*v5x;
    vy = C0*v0y + C1*v1y + C2*v2y + C3*v3y + C4*v4y + C5*v5y;   


    ax = M*vx;
    ay = M*vy;
    

    wfish[i].ax = ax;
    wfish[i].ay = ay;
  }

}

/* keep fish inside moving grid */
void
bounce_fish (fish_t* fish)
{
  while (fish->x < LEFT_WALL || fish->x > RIGHT_WALL) {
    if (fish->x < LEFT_WALL) {
      fish->x = fish->x + WALL_SEP;
    }
    if (fish->x > RIGHT_WALL) {
      fish->x = fish->x - WALL_SEP;
    }
  }
  while (fish->y < LEFT_WALL || fish->y > RIGHT_WALL) {
    if (fish->y < LEFT_WALL) {
      fish->y = fish->y + WALL_SEP;
    }
    if (fish->y > RIGHT_WALL) {
      fish->y = fish->y - WALL_SEP;
    }
  }
}

